# # https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment1/question3/main_interpolation.py
import math
import os
import typing

import matplotlib.pyplot
import torch

from interpolation_method import InterpolationMethod1d
from interpolations.cubic_spline import CubicSplineInterpolation
from interpolations.lagrange import LagrangeInterpolation
from interpolations.neville import NevilleInterpolation
from interpolations.newton import NewtonInterpolation


def _draw_function(f: typing.Callable[[float], float]):
    x = torch.linspace(-5, 5, 100)
    y = [f(v) for v in x]
    return matplotlib.pyplot.plot(x, y)[0]


def _get_interpolation_nodes(
        f: typing.Callable[[float], float],
        count: int):
    for x in torch.linspace(-5, 5, count):
        x = float(x)
        yield x, f(x)


def _test_interpolation_and_output(
        title: str,
        method: InterpolationMethod1d):
    def f(x: float):
        upper = math.sqrt(3) + math.sin(2 * x)
        lower = 3 + 4 * x ** 2
        return math.tan(math.cos(upper / lower))

    matplotlib.pyplot.rcParams['font.sans-serif'] = ['SimHei']
    matplotlib.pyplot.rcParams['axes.unicode_minus'] = False
    matplotlib.pyplot.figure(figsize=(12, 8), dpi=200)

    matplotlib.pyplot.subplot(2, 3, 6)
    matplotlib.pyplot.title(f"{title}（3 5 7 9 11）")
    matplotlib.pyplot.ylim(0.7, 1.9)
    _draw_function(f)

    for i, node_count in enumerate([3, 5, 7, 9, 11], 1):
        print(f"正在针对 {node_count} 个插值节点进行{title}，可能需要一定的时间……")
        nodes = list(_get_interpolation_nodes(f, node_count))
        result = method.interpolate(nodes)

        matplotlib.pyplot.subplot(2, 3, i)
        matplotlib.pyplot.title(f"{title}（{node_count}）")
        matplotlib.pyplot.ylim(0.7, 1.9)
        f_path = _draw_function(f)
        nodes_path = matplotlib.pyplot.scatter([n[0] for n in nodes], [n[1] for n in nodes])
        result_path = _draw_function(result)
        matplotlib.pyplot.legend(
            [f_path, nodes_path, result_path], ["原函数", "插值节点", "插值函数"], loc="lower right")

        matplotlib.pyplot.subplot(2, 3, 6)
        matplotlib.pyplot.scatter([n[0] for n in nodes], [n[1] for n in nodes])
        _draw_function(result)

    matplotlib.pyplot.savefig(f"./outputs/{title}.png")


def _main():
    _test_interpolation_and_output("拉格朗日插值", LagrangeInterpolation())
    _test_interpolation_and_output("牛顿插值", NewtonInterpolation())
    _test_interpolation_and_output("纳维尔插值", NevilleInterpolation())
    _test_interpolation_and_output("三次条样插值", CubicSplineInterpolation())
    print(f"运行完成，结果已保存在 {os.path.abspath('./outputs')}")


if __name__ == "__main__":
    _main()
